// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/type/simple_function_api.h>
#include <pollux/type/type.h>
#include <pollux/vector/complex_vector.h>

namespace kumo::pollux::test {

template <typename TKEY, typename TVALUE>
class MapBuilder {
 public:
  using pair = std::pair<TKEY, std::optional<TVALUE>>;
  using row = std::vector<pair>;
  using rows = std::vector<std::optional<row>>;

  static VectorPtr create(memory::MemoryPool& pool, rows rows) {
    const auto type = CppToType<Map<TKEY, TVALUE>>::create();
    auto items = 0;
    for (auto& row : rows) {
      if (row.has_value()) {
        items += row->size();
      }
    }

    BufferPtr nulls = allocate_nulls(rows.size(), &pool);
    auto* nullsPtr = nulls->asMutable<uint64_t>();
    size_t nullCount = 0;

    BufferPtr offsets = allocateOffsets(rows.size(), &pool);
    auto* offsetsPtr = offsets->asMutable<vector_size_t>();

    BufferPtr lengths = allocateSizes(rows.size(), &pool);
    auto* lengthsPtr = lengths->asMutable<vector_size_t>();

    BufferPtr keys = AlignedBuffer::allocate<TKEY>(items, &pool);
    auto* keysPtr = keys->asMutable<TKEY>();

    BufferPtr values = AlignedBuffer::allocate<TVALUE>(items, &pool);
    auto* valuesPtr = values->asMutable<TVALUE>();

    BufferPtr valueNulls = allocate_nulls(items, &pool);
    auto* valueNullsPtr = valueNulls->asMutable<uint64_t>();

    auto i = 0;
    auto offset = 0;

    for (auto& row : rows) {
      offsetsPtr[i] = offset;

      if (row.has_value()) {
        bits::clearNull(nullsPtr, i);
        lengthsPtr[i] = (*row).size();

        for (auto& pair : *row) {
          keysPtr[offset] = pair.first;
          if (pair.second.has_value()) {
            valuesPtr[offset] = *pair.second;
            bits::clearNull(valueNullsPtr, offset);
          } else {
            bits::set_null(valueNullsPtr, offset);
          }
          ++offset;
        }
      } else {
        nullCount++;
        bits::set_null(nullsPtr, i);
        lengthsPtr[i] = 0;
      }

      ++i;
    }

    offsetsPtr[i] = offset;

    return std::make_shared<MapVector>(
        &pool,
        type,
        nulls,
        rows.size(),
        offsets,
        lengths,
        std::make_shared<FlatVector<TKEY>>(
            &pool,
            type->childAt(0),
            BufferPtr(nullptr),
            items /*length*/,
            keys,
            std::vector<BufferPtr>()),
        std::make_shared<FlatVector<TVALUE>>(
            &pool,
            type->childAt(1),
            valueNulls,
            items /*length*/,
            values,
            std::vector<BufferPtr>()),
        nullCount);
  };
};

} // namespace kumo::pollux::test
